package simpledb.execution;

import simpledb.common.Type;
import simpledb.storage.*;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 能够对一组 IntField 字段进行聚合计算的类。
 *
 * 支持常见的聚合操作：MIN、MAX、SUM、AVG、COUNT 等；
 * 支持按字段分组（group by）或不进行分组；
 */
public class IntegerAggregator implements Aggregator {

    private static final long serialVersionUID = 1L;

    /**
     * 用于保存每个分组的聚合中间结果。
     */
    private static class AggInfo {
        int cnt;                    // 计数器，用于 COUNT 和 AVG
        int sum;                    // 求和值，用于 SUM 和 AVG
        int max = Integer.MIN_VALUE; // 最大值，用于 MAX
        int min = Integer.MAX_VALUE; // 最小值，用于 MIN
    }

    // 存储每个分组对应的聚合信息（Group By 字段 -> 聚合结果）
    private Map<Field, AggInfo> groupMap;

    // 分组字段在 Tuple 中的索引，若为 NO_GROUPING 表示不分组
    private int                 gbField;

    // 分组字段的数据类型
    private Type                gbFieldType;

    // 聚合字段在 Tuple 中的索引
    private int                 agField;

    // 使用的聚合操作类型（如 MIN、MAX、AVG 等）
    private Op                  op;

    // 默认字段（用于无分组时作为键）
    private Field               DEFAULT_FIELD = new StringField("Default", 10);

    // 缓存输出元组结构描述
    private TupleDesc           td;

    /**
     * 构造函数。
     *
     * @param gbfield 分组字段在 Tuple 中的索引（从 0 开始），若为 NO_GROUPING 表示不分组
     * @param gbfieldtype 分组字段的类型（如 Type.INT_TYPE），若不分组则为 null
     * @param afield 聚合字段在 Tuple 中的索引（从 0 开始）
     * @param what 使用的聚合操作（如 MIN、MAX、AVG 等）
     */
    public IntegerAggregator(int gbfield, Type gbfieldtype, int afield, Op what) {
        this.groupMap = new HashMap<>();
        this.gbField = gbfield;
        this.agField = afield;
        this.op = what;
        this.gbFieldType = gbfieldtype;
    }

    /**
     * 将一个元组合并进聚合器中。
     * 根据构造函数指定的分组方式，将该元组加入到相应的聚合组中。
     *
     * @param tup 包含聚合字段和分组字段的元组
     */
    public void mergeTupleIntoGroup(Tuple tup) {
        if (this.td == null) {
            buildTupleDesc(tup.getTupleDesc());
        }
        final IntField agFieldVal = (IntField) tup.getField(this.agField);
        final Field gbFieldVal = (this.gbField == NO_GROUPING) ? DEFAULT_FIELD : tup.getField(this.gbField);

        doAggregation(gbFieldVal, agFieldVal.getValue());
    }

    /**
     * 执行具体的聚合操作。
     *
     * @param key 分组键（即分组字段的值）
     * @param value 当前处理的聚合字段值
     */
    private void doAggregation(final Field key, final int value) {
        if (key != null) {
            AggInfo preInfo = this.groupMap.getOrDefault(key, new AggInfo());

            switch (this.op) {
                case MIN:
                    preInfo.min = Math.min(preInfo.min, value);
                    break;
                case MAX:
                    preInfo.max = Math.max(preInfo.max, value);
                    break;
                case AVG:
                case SUM_COUNT: // lab7 用
                    preInfo.sum += value;
                    preInfo.cnt += 1;
                    break;
                case SUM:
                    preInfo.sum += value;
                    break;
                case COUNT:
                    preInfo.cnt += 1;
                    break;
                default:
                    throw new UnsupportedOperationException("不支持的聚合操作：" + this.op);
            }

            this.groupMap.put(key, preInfo);
        }
    }

    /**
     * 根据聚合中间结果，获取最终聚合值。
     *
     * @param key 分组键
     * @return 对应的聚合结果值
     */
    private int parseValue(final Field key) {
        if (key != null && this.groupMap.containsKey(key)) {
            AggInfo preInfo = this.groupMap.get(key);

            switch (this.op) {
                case MIN:
                    return preInfo.min;
                case MAX:
                    return preInfo.max;
                case AVG:
                    return preInfo.sum / preInfo.cnt;
                case SUM:
                    return preInfo.sum;
                case COUNT:
                    return preInfo.cnt;
                case SUM_COUNT:
                    // 实际使用时会封装成两个字段返回，这里只返回其中一个
                    return preInfo.sum;
                default:
                    throw new UnsupportedOperationException("不支持的聚合操作：" + this.op);
            }
        }
        return 0;
    }

    /**
     * 构建输出元组的结构描述。
     *
     * 如果没有分组字段，则输出只有一个字段表示聚合结果；
     * 如果有分组字段，则第一个字段是分组字段值，第二个是聚合结果；
     *
     * @param originTd 输入元组的结构描述
     */
    public void buildTupleDesc(final TupleDesc originTd) {
        if (this.gbField == NO_GROUPING) {
            Type[] types = new Type[] { Type.INT_TYPE };
            String[] names = new String[] { originTd.getFieldName(this.agField) + "(" + this.op + ")" };
            this.td = new TupleDesc(types, names);
        } else {
            Type[] types = new Type[] { this.gbFieldType, Type.INT_TYPE };
            String[] names = new String[] { originTd.getFieldName(this.gbField),
                    originTd.getFieldName(this.agField) + "(" + this.op + ")" };
            this.td = new TupleDesc(types, names);
        }
    }

    /**
     * 创建一个迭代器，遍历所有分组的聚合结果。
     *
     * @return 迭代器对象：
     *         - 如果有分组，则返回 (groupVal, aggregateVal) 的元组；
     *         - 如果无分组，则返回单个字段 (aggregateVal)；
     *         - 具体值由构造时指定的聚合操作决定；
     */
    public OpIterator iterator() {
        final List<Tuple> tuples = new ArrayList<>();

        if (this.gbField != NO_GROUPING) {
            this.groupMap.forEach((key, info) -> {
                final Tuple tuple = new Tuple(this.td);
                tuple.setField(0, key);
                tuple.setField(1, new IntField(parseValue(key)));
                tuples.add(tuple);
            });
        } else {
            final Tuple tuple = new Tuple(this.td);
            tuple.setField(0, new IntField(parseValue(DEFAULT_FIELD)));
            tuples.add(tuple);
        }

        return new TupleIterator(this.td, tuples);
    }
}
